package com.cloudea.learn;

import com.cloudea.learn.exception.DimensionException;
import com.cloudea.learn.exception.TensorIndexException;

import java.util.*;

/**
 * 用于存储数据的结构
 * @author cloudea
 * @date 2021/12/28
 */
public class Tensor {
    public static final Random random = new Random(new Date().getTime());

    private final double[] data;
    private int[] shape ;
    private final Map<String, Object> attrs;

    public Tensor(int... dimensions){
        if(dimensions.length <= 0){
            throw new DimensionException(dimensions);
        }
        int numOfData = 1;
        for (var dimension : dimensions){
            // 某一维的维度可以是0，此时数据个数为0
            if(dimension < 0){
                throw new DimensionException(dimensions);
            }
            numOfData *= dimension;
        }

        this.shape = dimensions;
        this.data = new double[numOfData];
        this.attrs = new HashMap<>();
    }

    public Tensor(boolean normal, int... dimensions){
        this(dimensions);
        if(normal){
            fillGaussian();
        }
    }

    /**
     * 填充值
     * @param value 预填充的值
     */
    public void fill(double value) {
        Arrays.fill(data, value);
    }

    /**
     * 随机填充
     */
    public void fillGaussian(){
        for (int i = 0; i < data.length; i++){
            data[i] = random.nextGaussian();
        }
    }

    /**
     * 随机填充
     */
    public void fillUniform(){
        for (int i = 0; i < data.length; i++){
            data[i] = random.nextDouble();
        }
    }

    /**
     * 取数
     * @param index 要取数的位置，为了效率不会判断索引是否正确
     * @return 数
     */
    public double get(int... index){
        try {
            return data[indexToPos(index)];
        }catch (Exception e){
            throw new TensorIndexException(index, shape);
        }
    }

    /**
     * 置数
     * @param index 要置数的位置
     * @param number 要置的数
     */
    public void set(double number, int... index){
        try {
            data[indexToPos(index)] = number;
        }catch (Exception e){
            throw new TensorIndexException(index, shape);
        }
    }

    /**
     * 根据存储位置获取元素值
     * @param location 存储位置
     * @return 值
     */
    public double getByLocation(int location){
        return this.data[location];
    }

    /**
     * 据存储位置设置元素值
     * @param location 存储位置
     * @param value 值
     */
    public void setByLocation(int location, double value){
        this.data[location] = value;
    }

    /**
     * 对数组中每个元素做映射
     * @param lambda 映射函数
     */
    public void map(Lambda lambda){
        for (int i = 0; i < data.length; i++){
            data[i] = lambda.map(data[i]);
        }
    }

    /**
     * 复制一份tensor
     * @return 返回复制后的tensor
     */
    public Tensor copy(){
        Tensor tensor = new Tensor(shape);
        for (int i = 0; i < data.length; i++){
            tensor.data[i] = data[i];
        }
        return tensor;
    }

    /**
     * 返回形状
     * @return 形状
     */
    public int[] getShape(){
        return this.shape;
    }

    /**
     * 获取数据个数
     * @return 数据个数
     */
    public int getSize(){
        return this.data.length;
    }

    /**
     * 设置tensor的属性
     * @param name 属性名
     * @param value 属性值
     */
    public void setAttribute(String name, Object value){
        this.attrs.put(name, value);
    }

    /**
     * 获取tensor的属性
     * @param name 属性名
     * @return 属性值
     */
    public Object getAttribute(String name){
        return this.attrs.get(name);
    }

    /**
     * 更新数组的形状
     * @param newShape 新的形状
     */
    public void reshape(int... newShape){
        int m = 1;
        for (int i = 0; i < newShape.length; i++){
            m *= newShape[i];
        }
        if(m != this.data.length){
            throw new DimensionException(newShape);
        }
        this.shape = newShape;
    }

    @Override
    public String toString() {
        return "Tensor{" +
                "data=" + tensorToString() +
                ", shape=" + Arrays.toString(shape) +
                '}';
    }

    /**
     * 把索引转
     * @param index 索引
     * @return 数据的位置
     */
    private int indexToPos(int... index){
        int weight = 1;
        int position = 0;
        for (int i = index.length - 1; i >= 0 ; i--){
            position += index[i] * weight;
            weight *= shape[i];
        }
        return position;
    }

    /**
     * 把第l层转成字符串
     * @param index 索引，从0开始
     * @return 字符串表示
     */
    private String tensorToString(int... index){
        if(index.length == shape.length - 1){
            StringBuilder stringBuilder = new StringBuilder();
            stringBuilder.append("[");
            for (int i = 0; i < shape[index.length]; i++){
                int[] indexNew = new int[shape.length];
                for (int j = 0; j < shape.length - 1; j++){
                    indexNew[j] = index[j];
                }
                indexNew[shape.length - 1] = i;
                if(i != 0){
                    stringBuilder.append(" ");
                }
                double number = get(indexNew);
                stringBuilder.append(number);
            }
            stringBuilder.append("]");
            return stringBuilder.toString();
        }else {
            StringBuilder stringBuilder = new StringBuilder();
            stringBuilder.append("[");
            for (int i = 0; i < shape[index.length]; i++){
                int[] indexNew = new int[index.length + 1];
                for (int j = 0; j < index.length; j++){
                    indexNew[j] = index[j];
                }
                indexNew[indexNew.length - 1] = i;
                if(i != 0){
                    stringBuilder.append(",");
                }
                stringBuilder.append(tensorToString(indexNew));
            }
            stringBuilder.append("]");
            return stringBuilder.toString();
        }
    }

    /**
     * 一一映射接口
     */
    public static interface Lambda {
        double map(double map);
    }

}
